import os
import cv2
import numpy as np
from PIL import Image
import requests
from io import BytesIO
import time
import logging

# 本地数字图片路径
TEMPLATE_DIR = "/ragflow/api/apps/sdk/img/"  # 确保文件夹中有 0.gif ~ 9.gif



class ImageRecognizer:
    """
    图片识别类，用于从图片中识别数字。

    """
    def __init__(self,download_delay=0.5):
        self.templates = self.load_template_images()
        self.page_img_cache = {}  # 页面级图片缓存 {img_url: recognized_number}
        self.download_delay = download_delay  # 控制下载延迟


    def clear_cache(self):
        """清空当前页面的图片缓存"""
        self.page_img_cache.clear()

    def load_template_images(self):
        """加载本地 0~9 的模板图片"""
        templates = {}
        for num in range(10):
            img_path = os.path.join(TEMPLATE_DIR, f"{num}.gif")
            try:
                img = Image.open(img_path).convert("L")  # 转为灰度
                templates[num] = np.array(img)  # 转为 NumPy 数组供 OpenCV 使用
                logging.info(f"✅ 成功加载模板: {img_path}")
            except Exception as e:
                logging.info(f"❌ 加载模板失败: {img_path} - {e}")
        return templates

    def download_target_image(self, url):
        """下载图片并自动添加延迟"""
        try:
            response = requests.get(url, timeout=5)
            response.raise_for_status()
            
            # 下载完成后添加延迟
            if self.download_delay > 0:
                time.sleep(self.download_delay)
                
            img = Image.open(BytesIO(response.content)).convert("L")
            return np.array(img)
        except Exception as e:
            logging.info(f"❌ 下载目标图片失败: {e}")
            return None

    def recognize_image(self, img_url):
        """识别图片中的数字，返回(识别结果, 是否来自缓存)"""
        # 检查缓存
        if img_url in self.page_img_cache:
            return self.page_img_cache[img_url],True
        
        if not self.templates:
            logging.info("❌ 无有效模板，请检查 img/ 目录！")
            return "[识别失败]"
        
        target_img = self.download_target_image(img_url)
        if target_img is None:
            return "[图片下载失败]"
        
        # 模板匹配
        best_num, best_score = -1, -np.inf
        for num, template in self.templates.items():
            resized_template = cv2.resize(template, (target_img.shape[1], target_img.shape[0]))
            result = cv2.matchTemplate(target_img, resized_template, cv2.TM_CCOEFF_NORMED)
            _, max_val, _, _ = cv2.minMaxLoc(result)
            if max_val > best_score:
                best_num, best_score = num, max_val
        
        if best_score > 0.7:  # 设置相似度阈值
            result = str(best_num)
            self.page_img_cache[img_url] = result  # 存入缓存
            return result,False
        else:
            logging.info(f"⚠️ 低置信度匹配: 数字 {best_num} (相似度: {best_score:.2f})")
            return "[低置信度]",False